#!/usr/bin/env python
# encoding: utf-8
'''
@file: video2imgs.py
@author: wangjianong
@time: 2020/3/27 17:06
@desc:
'''

import cv2
import os
from multiprocessing import Process,Pool
import argparse

def video2imgs(video_path,save_folder,save_prefix=None,interval=5):
    if os.path.isdir(video_path):
        list_videos = os.listdir(video_path)
        list_videos = [os.path.join(video_path,p) for p in list_videos]
        # list_videos = list(filter(lambda p:os.path.isfile,list_videos))
        for video_path in list_videos:
            video2imgs(video_path,save_folder,interval=interval)
    else:
        if save_prefix is None:
            save_prefix = os.path.splitext(os.path.basename(video_path))[0]
        os.makedirs(save_folder,exist_ok=True)
        cap = cv2.VideoCapture(video_path)
        assert cap.isOpened(),f'{video_path} can not open!'
        frame_id = -1
        while True:
            ret,frame = cap.read()
            if frame is None:
                break
            frame_id += 1
            if frame_id %interval !=0:
                continue
            frame_name = save_prefix +'_' + str(frame_id).zfill(6)+'.png'
            save_path = os.path.join(save_folder,frame_name)
            cv2.imwrite(save_path,frame)
        cap.release()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('src_path',type=str)
    parser.add_argument('--dst_path',default='imgs/',type=str)
    parser.add_argument('--interval',type=int,default=5)
    parser.add_argument('--prefix',type=str,default=None)
    args = parser.parse_args()

    video2imgs(args.src_path,args.dst_path,args.prefix,args.interval)

